class PasswordProtectedFile
  DEFAULT_ITERATIONS = 20_000
  ITERATION_BYTES = 8
  KEY_BYTES = 32
  SALT_BYTES = 16
  IV_BYTES = 16
  HASH_FUNCTION = OpenSSL::Digest::SHA256

  private_class_method :new

  def self.open(file_name, password)
    assert_valid_filename(file_name)
    assert_file_exists(file_name)
    assert_valid_password(password)
    new(file_name, password).tap { |o| o.__send__(:open_existing) }
  end

  def self.create(file_name, password, data = '')
    assert_valid_filename(file_name)
    assert_file_does_not_exist(file_name)
    assert_valid_password(password)
    new(file_name, password).tap { |o| o.__send__(:create_new, data) }
  end

  def data
    @data.clone
  end

  def data=(new_data)
    assert_valid_data(new_data)
    @data = new_data.clone
    write_file
  end

  private

  def initialize(file_name, password)
    @file_name = file_name
    @password = password
  end

  def create_new(data)
    @data = data
    @iterations = DEFAULT_ITERATIONS
    write_file
  end

  def open_existing
    load_file(@file_name, @password)
  end

  def load_file(file_name, password)
    file = File.binread(file_name).chars
    pw_salt = file.shift(SALT_BYTES).join
    pw_hash = file.shift(KEY_BYTES).join
    aes_pw_salt = file.shift(SALT_BYTES).join
    aes_iv = file.shift(IV_BYTES).join
    @iterations = file.shift(ITERATION_BYTES).join.to_i
    encrypted_data = file.join

    hashed_pw = pbkdf2(password, pw_salt)
    fail IncorrectPasswordError unless hashed_pw == pw_hash

    aes_key = pbkdf2(password, aes_pw_salt)
    cipher = new_aes256(:decrypt, aes_key, aes_iv)
    @data = (cipher.update(encrypted_data) + cipher.final)[0..-2]
  end

  def write_file
    pw_salt = new_salt
    pw_hash = pbkdf2(@password, pw_salt)
    aes_pw_salt = new_salt
    aes_iv = SecureRandom.random_bytes(IV_BYTES)
    aes_key = pbkdf2(@password, aes_pw_salt)

    cipher = new_aes256(:encrypt, aes_key, aes_iv)
    encrypted_data = cipher.update(@data + "\0") + cipher.final

    File.open(@file_name, 'w') do |f|
      f.print(pw_salt)
      f.print(pw_hash)
      f.print(aes_pw_salt)
      f.print(aes_iv)
      f.print(@iterations.to_s.rjust(ITERATION_BYTES))
      f.print(encrypted_data)
    end
    true
  end

  def pbkdf2(password, salt)
    OpenSSL::PKCS5::pbkdf2_hmac(password, salt, @iterations, KEY_BYTES, HASH_FUNCTION.new)
  end

  def new_aes256(mode, key, iv)
    OpenSSL::Cipher::AES256.new(:CBC).tap do |c|
      c.__send__(mode)
      c.key = key
      c.iv = iv
    end
  end

  def new_salt
    SecureRandom.random_bytes(SALT_BYTES)
  end

  def self.assert_valid_filename(file_name)
    unless file_name.is_a?(String) && !file_name.empty? && Dir.exist?(File.dirname(file_name))
      fail InvalidFilenameError.new("Invalid filename given: #{file_name.inspect}")
    end
  end

  def self.assert_file_exists(file_name)
    fail FileNotFoundError.new("File #{file_name} not found") unless File.exist?(file_name)
  end

  def self.assert_file_does_not_exist(file_name)
    fail FilenameNotAvailableError.new("File #{file_name} already exists") if File.exist?(file_name)
  end

  def self.assert_valid_password(password)
    fail InvalidPasswordError.new("Invalid password given: #{password.inspect}") unless password.is_a?(String) && password.length > 0
  end

  def assert_valid_data(d)
    fail InvalidDataError.new("Expected String, got #{d.class}") unless d.instance_of?(String)
    fail InvalidStringEncodingError.new("Expected utf-8, got #{d.encoding}") unless d.encoding == Encoding::UTF_8
  end

  class PasswordProtectedFileError < StandardError; end
  class IncorrectPasswordError < PasswordProtectedFileError; end
  class InvalidPasswordError < PasswordProtectedFileError; end
  class InvalidDataError < PasswordProtectedFileError; end
  class InvalidStringEncodingError < PasswordProtectedFileError; end
  class InvalidFilenameError < PasswordProtectedFileError; end
  class FileNotFoundError < PasswordProtectedFileError; end
  class FilenameNotAvailableError < PasswordProtectedFileError; end
end
